"""
实验结果数据库CRUD操作
"""
from sqlalchemy.orm import Session
from sqlalchemy import and_
from typing import List, Optional, Dict, Any
from datetime import datetime

from .models import (
    ExperimentRun,
    ExperimentResult,
    ExperimentDataCategory,
    ExperimentDataSeries,
    ExperimentDataScalar
)
from ..parameters.models import Project


# ==================== 实验运行管理 ====================

def create_experiment_run(db: Session, project_id: int, name: Optional[str] = None, 
                         description: Optional[str] = None, 
                         snapshot_file_path: Optional[str] = None,
                         created_at: Optional[Any] = None) -> ExperimentRun:
    """创建实验运行记录
    
    Args:
        db: 数据库会话
        project_id: 项目ID
        name: 实验名称
        description: 实验描述
        snapshot_file_path: 快照文件路径
        created_at: 创建时间（可选，如果不提供则使用当前时间）
    """
    # 如果提供了created_at，先解析为datetime对象
    parsed_created_at = None
    if created_at is not None:
        # 如果传入的是字符串，解析为datetime
        if isinstance(created_at, str):
            try:
                # 处理ISO格式字符串，支持带时区和不带时区的格式
                if 'Z' in created_at:
                    parsed_created_at = datetime.fromisoformat(created_at.replace('Z', '+00:00'))
                else:
                    parsed_created_at = datetime.fromisoformat(created_at)
                # 如果是带时区的datetime，转换为本地naive datetime
                if parsed_created_at.tzinfo is not None:
                    parsed_created_at = parsed_created_at.replace(tzinfo=None)
            except ValueError:
                # 如果解析失败，使用None（让数据库使用默认值）
                parsed_created_at = None
        elif isinstance(created_at, datetime):
            parsed_created_at = created_at
            # 如果是带时区的datetime，转换为本地naive datetime
            if parsed_created_at.tzinfo is not None:
                parsed_created_at = parsed_created_at.replace(tzinfo=None)
    
    # 在创建对象时直接传入created_at（如果提供了的话）
    experiment_kwargs = {
        "project_id": project_id,
        "name": name,
        "description": description,
        "snapshot_file_path": snapshot_file_path
    }
    if parsed_created_at is not None:
        experiment_kwargs["created_at"] = parsed_created_at
    
    db_experiment = ExperimentRun(**experiment_kwargs)
    
    db.add(db_experiment)
    db.commit()
    # 使用 flush 而不是 refresh，避免重新从数据库加载（可能覆盖我们设置的值）
    db.flush()
    db.refresh(db_experiment)
    return db_experiment


def get_experiment_run(db: Session, experiment_run_id: int) -> Optional[ExperimentRun]:
    """根据ID获取实验运行记录"""
    return db.query(ExperimentRun).filter(ExperimentRun.id == experiment_run_id).first()


def get_experiment_runs_by_project(db: Session, project_id: int, 
                                   skip: int = 0, limit: int = 100) -> List[ExperimentRun]:
    """获取项目的实验运行记录列表"""
    return db.query(ExperimentRun).filter(
        ExperimentRun.project_id == project_id
    ).order_by(ExperimentRun.created_at.desc()).offset(skip).limit(limit).all()


def get_experiment_runs_by_project_name_en(db: Session, project_name_en: str,
                                           skip: int = 0, limit: int = 100) -> List[ExperimentRun]:
    """根据项目英文名称获取实验运行记录列表"""
    project = db.query(Project).filter(Project.name_en == project_name_en).first()
    if not project:
        return []
    return get_experiment_runs_by_project(db, project.id, skip, limit)


def update_experiment_run(db: Session, experiment_run_id: int, 
                         name: Optional[str] = None,
                         description: Optional[str] = None) -> Optional[ExperimentRun]:
    """更新实验运行记录"""
    db_experiment = get_experiment_run(db, experiment_run_id)
    if db_experiment:
        if name is not None:
            db_experiment.name = name
        if description is not None:
            db_experiment.description = description
        db.commit()
        db.refresh(db_experiment)
    return db_experiment


def delete_experiment_run(db: Session, experiment_run_id: int) -> bool:
    """删除实验运行记录"""
    db_experiment = get_experiment_run(db, experiment_run_id)
    if db_experiment:
        db.delete(db_experiment)
        db.commit()
        return True
    return False


# ==================== 实验结果文件管理 ====================

def create_experiment_result(db: Session, experiment_run_id: int, name: str,
                             file_type: str, file_path: str,
                             file_size: Optional[int] = None,
                             description: Optional[str] = None) -> ExperimentResult:
    """创建实验结果文件记录"""
    db_result = ExperimentResult(
        experiment_run_id=experiment_run_id,
        name=name,
        file_type=file_type,
        file_path=file_path,
        file_size=file_size,
        description=description
    )
    db.add(db_result)
    db.commit()
    db.refresh(db_result)
    return db_result


def get_experiment_result(db: Session, result_id: int) -> Optional[ExperimentResult]:
    """根据ID获取实验结果文件记录"""
    return db.query(ExperimentResult).filter(ExperimentResult.id == result_id).first()


def get_experiment_results(db: Session, experiment_run_id: int) -> List[ExperimentResult]:
    """获取实验的所有结果文件"""
    return db.query(ExperimentResult).filter(
        ExperimentResult.experiment_run_id == experiment_run_id
    ).order_by(ExperimentResult.created_at).all()


def get_experiment_results_by_type(db: Session, experiment_run_id: int, 
                                   file_type: str) -> List[ExperimentResult]:
    """根据文件类型获取实验结果文件"""
    return db.query(ExperimentResult).filter(
        and_(
            ExperimentResult.experiment_run_id == experiment_run_id,
            ExperimentResult.file_type == file_type
        )
    ).order_by(ExperimentResult.created_at).all()


def update_experiment_result(db: Session, result_id: int,
                            name: Optional[str] = None,
                            description: Optional[str] = None) -> Optional[ExperimentResult]:
    """更新实验结果文件记录"""
    db_result = get_experiment_result(db, result_id)
    if db_result:
        if name is not None:
            db_result.name = name
        if description is not None:
            db_result.description = description
        db.commit()
        db.refresh(db_result)
    return db_result


def delete_experiment_result(db: Session, result_id: int) -> bool:
    """删除实验结果文件记录"""
    db_result = get_experiment_result(db, result_id)
    if db_result:
        db.delete(db_result)
        db.commit()
        return True
    return False


# ==================== 数据分类管理 ====================

def create_experiment_data_category(db: Session, experiment_run_id: int, name: str,
                                   description: Optional[str] = None,
                                   sort_order: int = 0) -> ExperimentDataCategory:
    """创建数据分类"""
    db_category = ExperimentDataCategory(
        experiment_run_id=experiment_run_id,
        name=name,
        description=description,
        sort_order=sort_order
    )
    db.add(db_category)
    db.commit()
    db.refresh(db_category)
    return db_category


def get_experiment_data_category(db: Session, category_id: int) -> Optional[ExperimentDataCategory]:
    """根据ID获取数据分类"""
    return db.query(ExperimentDataCategory).filter(ExperimentDataCategory.id == category_id).first()


def get_experiment_data_category_by_name(db: Session, experiment_run_id: int, 
                                         name: str) -> Optional[ExperimentDataCategory]:
    """根据名称获取数据分类"""
    return db.query(ExperimentDataCategory).filter(
        and_(
            ExperimentDataCategory.experiment_run_id == experiment_run_id,
            ExperimentDataCategory.name == name
        )
    ).first()


def get_experiment_data_categories(db: Session, experiment_run_id: int) -> List[ExperimentDataCategory]:
    """获取实验的所有数据分类"""
    return db.query(ExperimentDataCategory).filter(
        ExperimentDataCategory.experiment_run_id == experiment_run_id
    ).order_by(ExperimentDataCategory.sort_order, ExperimentDataCategory.created_at).all()


def update_experiment_data_category(db: Session, category_id: int,
                                   name: Optional[str] = None,
                                   description: Optional[str] = None,
                                   sort_order: Optional[int] = None) -> Optional[ExperimentDataCategory]:
    """更新数据分类"""
    db_category = get_experiment_data_category(db, category_id)
    if db_category:
        if name is not None:
            db_category.name = name
        if description is not None:
            db_category.description = description
        if sort_order is not None:
            db_category.sort_order = sort_order
        db.commit()
        db.refresh(db_category)
    return db_category


def delete_experiment_data_category(db: Session, category_id: int) -> bool:
    """删除数据分类
    
    删除分类时，关联数据的category_id会自动设置为NULL（通过外键约束SET NULL）。
    上层代码需要在删除前手动迁移数据到默认分类。
    """
    db_category = get_experiment_data_category(db, category_id)
    if db_category:
        db.delete(db_category)
        db.commit()
        return True
    return False


# ==================== 序列数据管理 ====================

def create_experiment_data_series(db: Session, experiment_run_id: int, category_id: int,
                                  name: str, parquet_path: str, data_length: int,
                                  data_group: Optional[str] = None,
                                  unit: Optional[str] = None,
                                  description: Optional[str] = None,
                                  index_column_name: Optional[str] = None) -> ExperimentDataSeries:
    """创建序列数据"""
    db_series = ExperimentDataSeries(
        experiment_run_id=experiment_run_id,
        category_id=category_id,
        name=name,
        data_group=data_group,
        parquet_path=parquet_path,
        unit=unit,
        description=description,
        index_column_name=index_column_name or "index",
        data_length=data_length
    )
    db.add(db_series)
    db.commit()
    db.refresh(db_series)
    return db_series


def get_experiment_data_series(db: Session, series_id: int) -> Optional[ExperimentDataSeries]:
    """根据ID获取序列数据"""
    return db.query(ExperimentDataSeries).filter(ExperimentDataSeries.id == series_id).first()


def get_experiment_data_series_by_name(db: Session, experiment_run_id: int, 
                                       category_id: int, name: str) -> Optional[ExperimentDataSeries]:
    """根据名称获取序列数据"""
    return db.query(ExperimentDataSeries).filter(
        and_(
            ExperimentDataSeries.experiment_run_id == experiment_run_id,
            ExperimentDataSeries.category_id == category_id,
            ExperimentDataSeries.name == name
        )
    ).first()


def get_experiment_data_series_by_category(db: Session, category_id: int) -> List[ExperimentDataSeries]:
    """获取分类下的所有序列数据"""
    return db.query(ExperimentDataSeries).filter(
        ExperimentDataSeries.category_id == category_id
    ).order_by(ExperimentDataSeries.name).all()


def get_experiment_data_series_by_experiment(db: Session, experiment_run_id: int) -> List[ExperimentDataSeries]:
    """获取实验的所有序列数据"""
    return db.query(ExperimentDataSeries).filter(
        ExperimentDataSeries.experiment_run_id == experiment_run_id
    ).order_by(ExperimentDataSeries.category_id, ExperimentDataSeries.name).all()


def get_experiment_data_series_by_group(db: Session, experiment_run_id: int,
                                       data_group: str) -> List[ExperimentDataSeries]:
    """根据数据分组获取序列数据"""
    return db.query(ExperimentDataSeries).filter(
        and_(
            ExperimentDataSeries.experiment_run_id == experiment_run_id,
            ExperimentDataSeries.data_group == data_group
        )
    ).order_by(ExperimentDataSeries.name).all()


def update_experiment_data_series(db: Session, series_id: int,
                                  name: Optional[str] = None,
                                  data_group: Optional[str] = None,
                                  unit: Optional[str] = None,
                                  description: Optional[str] = None,
                                  data_length: Optional[int] = None) -> Optional[ExperimentDataSeries]:
    """更新序列数据"""
    db_series = get_experiment_data_series(db, series_id)
    if db_series:
        if name is not None:
            db_series.name = name
        if data_group is not None:
            db_series.data_group = data_group
        if unit is not None:
            db_series.unit = unit
        if description is not None:
            db_series.description = description
        if data_length is not None:
            db_series.data_length = data_length
        db.commit()
        db.refresh(db_series)
    return db_series


def delete_experiment_data_series(db: Session, series_id: int) -> bool:
    """删除序列数据"""
    db_series = get_experiment_data_series(db, series_id)
    if db_series:
        db.delete(db_series)
        db.commit()
        return True
    return False


# ==================== 单值数据管理 ====================

def create_experiment_data_scalar(db: Session, experiment_run_id: int, category_id: int,
                                 name: str, value: str,
                                 data_group: Optional[str] = None,
                                 unit: Optional[str] = None,
                                 description: Optional[str] = None) -> ExperimentDataScalar:
    """创建单值数据"""
    db_scalar = ExperimentDataScalar(
        experiment_run_id=experiment_run_id,
        category_id=category_id,
        name=name,
        data_group=data_group,
        value=value,
        unit=unit,
        description=description
    )
    db.add(db_scalar)
    db.commit()
    db.refresh(db_scalar)
    return db_scalar


def get_experiment_data_scalar(db: Session, scalar_id: int) -> Optional[ExperimentDataScalar]:
    """根据ID获取单值数据"""
    return db.query(ExperimentDataScalar).filter(ExperimentDataScalar.id == scalar_id).first()


def get_experiment_data_scalar_by_name(db: Session, experiment_run_id: int,
                                       category_id: int, name: str) -> Optional[ExperimentDataScalar]:
    """根据名称获取单值数据"""
    return db.query(ExperimentDataScalar).filter(
        and_(
            ExperimentDataScalar.experiment_run_id == experiment_run_id,
            ExperimentDataScalar.category_id == category_id,
            ExperimentDataScalar.name == name
        )
    ).first()


def get_experiment_data_scalars_by_category(db: Session, category_id: int) -> List[ExperimentDataScalar]:
    """获取分类下的所有单值数据"""
    return db.query(ExperimentDataScalar).filter(
        ExperimentDataScalar.category_id == category_id
    ).order_by(ExperimentDataScalar.name).all()


def get_experiment_data_scalars_by_experiment(db: Session, experiment_run_id: int) -> List[ExperimentDataScalar]:
    """获取实验的所有单值数据"""
    return db.query(ExperimentDataScalar).filter(
        ExperimentDataScalar.experiment_run_id == experiment_run_id
    ).order_by(ExperimentDataScalar.category_id, ExperimentDataScalar.name).all()


def get_experiment_data_scalars_by_group(db: Session, experiment_run_id: int,
                                         data_group: str) -> List[ExperimentDataScalar]:
    """根据数据分组获取单值数据"""
    return db.query(ExperimentDataScalar).filter(
        and_(
            ExperimentDataScalar.experiment_run_id == experiment_run_id,
            ExperimentDataScalar.data_group == data_group
        )
    ).order_by(ExperimentDataScalar.name).all()


def update_experiment_data_scalar(db: Session, scalar_id: int,
                                  name: Optional[str] = None,
                                  data_group: Optional[str] = None,
                                  value: Optional[str] = None,
                                  unit: Optional[str] = None,
                                  description: Optional[str] = None) -> Optional[ExperimentDataScalar]:
    """更新单值数据"""
    db_scalar = get_experiment_data_scalar(db, scalar_id)
    if db_scalar:
        if name is not None:
            db_scalar.name = name
        if data_group is not None:
            db_scalar.data_group = data_group
        if value is not None:
            db_scalar.value = value
        if unit is not None:
            db_scalar.unit = unit
        if description is not None:
            db_scalar.description = description
        db.commit()
        db.refresh(db_scalar)
    return db_scalar


def delete_experiment_data_scalar(db: Session, scalar_id: int) -> bool:
    """删除单值数据"""
    db_scalar = get_experiment_data_scalar(db, scalar_id)
    if db_scalar:
        db.delete(db_scalar)
        db.commit()
        return True
    return False


# ==================== 实验完整数据查询 ====================

def get_experiment_full_data(db: Session, experiment_run_id: int) -> Dict[str, Any]:
    """获取实验的完整数据（包括分类、序列数据、单值数据、结果文件）"""
    experiment = get_experiment_run(db, experiment_run_id)
    if not experiment:
        return {}
    
    categories = get_experiment_data_categories(db, experiment_run_id)
    results = get_experiment_results(db, experiment_run_id)
    
    category_dict = {}
    for category in categories:
        series_data = get_experiment_data_series_by_category(db, category.id)
        scalar_data = get_experiment_data_scalars_by_category(db, category.id)
        
        category_dict[category.id] = {
            "id": category.id,
            "name": category.name,
            "description": category.description,
            "sort_order": category.sort_order,
            "created_at": category.created_at,
            "series_data": [
                {
                    "id": s.id,
                    "name": s.name,
                    "data_group": s.data_group,
                    "parquet_path": s.parquet_path,
                    "unit": s.unit,
                    "description": s.description,
                    "index_column_name": s.index_column_name,
                    "data_length": s.data_length,
                    "created_at": s.created_at
                }
                for s in series_data
            ],
            "scalar_data": [
                {
                    "id": s.id,
                    "name": s.name,
                    "data_group": s.data_group,
                    "value": s.value,
                    "unit": s.unit,
                    "description": s.description,
                    "created_at": s.created_at
                }
                for s in scalar_data
            ]
        }
    
    return {
        "id": experiment.id,
        "project_id": experiment.project_id,
        "name": experiment.name,
        "description": experiment.description,
        "snapshot_file_path": experiment.snapshot_file_path,
        "created_at": experiment.created_at,
        "results": [
            {
                "id": r.id,
                "name": r.name,
                "file_type": r.file_type,
                "file_path": r.file_path,
                "file_size": r.file_size,
                "description": r.description,
                "created_at": r.created_at
            }
            for r in results
        ],
        "categories": list(category_dict.values())
    }

